"""
供应链需求预测模块
提供多种预测算法和集成预测功能
"""

import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import warnings
warnings.filterwarnings('ignore')

class DemandForecaster:
    """需求预测核心类"""
    
    def __init__(self):
        self.forecast_methods = {
            'moving_average': self._moving_average_forecast,
            'exponential_smoothing': self._exponential_smoothing_forecast,
            'seasonal_decompose': self._seasonal_decompose_forecast,
            'trend_analysis': self._trend_analysis_forecast
        }
        self.accuracy_metrics = {}
    
    def prepare_demand_data(self, orders_df, sku_col='sku', date_col='order_date', qty_col='quantity'):
        """
        准备需求数据用于预测
        
        Args:
            orders_df: 订单数据DataFrame
            sku_col: SKU列名
            date_col: 日期列名
            qty_col: 数量列名
            
        Returns:
            dict: 按SKU分组的需求时间序列数据
        """
        demand_data = {}
        
        # 按SKU分组
        for sku, group in orders_df.groupby(sku_col):
            # 按日期聚合需求
            daily_demand = group.groupby(date_col)[qty_col].sum().reset_index()
            daily_demand[date_col] = pd.to_datetime(daily_demand[date_col])
            daily_demand = daily_demand.sort_values(date_col)
            daily_demand.set_index(date_col, inplace=True)
            
            # 重采样为日需求，填充缺失日期
            daily_demand = daily_demand.resample('D').sum().fillna(0)
            
            demand_data[sku] = daily_demand
        
        return demand_data
    
    def _moving_average_forecast(self, data, window=7, horizon=30):
        """移动平均预测"""
        if len(data) < window:
            # 数据不足时使用简单平均
            forecast = [data['quantity'].mean()] * horizon
        else:
            # 使用移动平均
            ma_values = data['quantity'].rolling(window=window).mean()
            last_ma = ma_values.dropna().iloc[-1] if not ma_values.dropna().empty else data['quantity'].mean()
            forecast = [max(0, last_ma)] * horizon
        
        return pd.Series(forecast, index=pd.date_range(
            start=data.index[-1] + timedelta(days=1),
            periods=horizon,
            freq='D'
        ))
    
    def _exponential_smoothing_forecast(self, data, alpha=0.3, horizon=30):
        """指数平滑预测"""
        if len(data) == 0:
            return pd.Series([0] * horizon)
        
        # 计算指数平滑
        smoothed = data['quantity'].ewm(alpha=alpha, adjust=False).mean()
        last_smoothed = smoothed.iloc[-1] if not smoothed.empty else data['quantity'].mean()
        
        forecast = [max(0, last_smoothed)] * horizon
        
        return pd.Series(forecast, index=pd.date_range(
            start=data.index[-1] + timedelta(days=1),
            periods=horizon,
            freq='D'
        ))
    
    def _seasonal_decompose_forecast(self, data, season_period=7, horizon=30):
        """季节性分解预测"""
        if len(data) < season_period * 2:
            # 数据不足，回退到移动平均
            return self._moving_average_forecast(data, horizon=horizon)
        
        # 确保索引是DatetimeIndex类型
        if not isinstance(data.index, pd.DatetimeIndex):
            data.index = pd.to_datetime(data.index)
        
        # 计算季节性因子
        seasonal_avg = data['quantity'].groupby(data.index.dayofweek).mean()
        overall_avg = data['quantity'].mean()
        seasonal_factors = seasonal_avg / overall_avg
        
        # 计算趋势（使用移动平均平滑）
        trend = data['quantity'].rolling(window=7).mean()
        last_trend = trend.dropna().iloc[-1] if not trend.dropna().empty else overall_avg
        
        # 生成预测
        forecast_dates = pd.date_range(
            start=data.index[-1] + timedelta(days=1),
            periods=horizon,
            freq='D'
        )
        
        forecast = []
        for date in forecast_dates:
            seasonal_factor = seasonal_factors[date.dayofweek]
            forecast_value = max(0, last_trend * seasonal_factor)
            forecast.append(forecast_value)
        
        return pd.Series(forecast, index=forecast_dates)
    
    def _trend_analysis_forecast(self, data, horizon=30):
        """趋势分析预测"""
        if len(data) < 14:
            # 数据不足，回退到移动平均
            return self._moving_average_forecast(data, horizon=horizon)
        
        # 确保索引是DatetimeIndex类型
        if not isinstance(data.index, pd.DatetimeIndex):
            data.index = pd.to_datetime(data.index)
        
        # 计算线性趋势
        x = np.arange(len(data))
        y = data['quantity'].values
        
        # 移除异常值
        q1, q3 = np.percentile(y, [25, 75])
        iqr = q3 - q1
        mask = (y >= q1 - 1.5 * iqr) & (y <= q3 + 1.5 * iqr)
        
        if np.sum(mask) > 1:
            x_filtered = x[mask]
            y_filtered = y[mask]
            
            # 线性回归
            coeffs = np.polyfit(x_filtered, y_filtered, 1)
            trend_slope = max(coeffs[0], -coeffs[0] * 0.1)  # 限制下降趋势
        else:
            trend_slope = 0
        
        # 生成预测
        last_value = data['quantity'].iloc[-1]
        forecast_dates = pd.date_range(
            start=data.index[-1] + timedelta(days=1),
            periods=horizon,
            freq='D'
        )
        
        forecast = []
        for i, date in enumerate(forecast_dates):
            trend_value = max(0, last_value + trend_slope * (i + 1))
            forecast.append(trend_value)
        
        return pd.Series(forecast, index=forecast_dates)
    
    def ensemble_forecast(self, data, methods=None, weights=None, horizon=30):
        """
        集成预测（多种方法加权平均）
        
        Args:
            data: 历史需求数据
            methods: 使用的方法列表
            weights: 各方法的权重
            horizon: 预测期长度
            
        Returns:
            pd.Series: 集成预测结果
        """
        if methods is None:
            methods = ['moving_average', 'exponential_smoothing', 'seasonal_decompose']
        
        if weights is None:
            weights = [1/len(methods)] * len(methods)
        
        # 获取各方法的预测结果
        forecasts = []
        for method in methods:
            if method in self.forecast_methods:
                forecast = self.forecast_methods[method](data, horizon=horizon)
                forecasts.append(forecast)
        
        if not forecasts:
            return self._moving_average_forecast(data, horizon=horizon)
        
        # 加权平均
        ensemble_forecast = pd.Series(0, index=forecasts[0].index)
        for forecast, weight in zip(forecasts, weights):
            ensemble_forecast += forecast * weight
        
        return ensemble_forecast
    
    def calculate_forecast_accuracy(self, actual, predicted, metrics=None):
        """
        计算预测准确性指标
        
        Args:
            actual: 实际值
            predicted: 预测值
            metrics: 要计算的指标列表
            
        Returns:
            dict: 各指标的计算结果
        """
        if metrics is None:
            metrics = ['mae', 'mse', 'rmse', 'mape', 'wape']
        
        actual = np.array(actual)
        predicted = np.array(predicted)
        
        # 移除零值和负值
        mask = (actual > 0) & (predicted >= 0)
        actual_filtered = actual[mask]
        predicted_filtered = predicted[mask]
        
        if len(actual_filtered) == 0:
            return {metric: np.inf for metric in metrics}
        
        results = {}
        
        # 平均绝对误差
        if 'mae' in metrics:
            results['mae'] = np.mean(np.abs(actual_filtered - predicted_filtered))
        
        # 均方误差
        if 'mse' in metrics:
            results['mse'] = np.mean((actual_filtered - predicted_filtered) ** 2)
        
        # 均方根误差
        if 'rmse' in metrics:
            results['rmse'] = np.sqrt(np.mean((actual_filtered - predicted_filtered) ** 2))
        
        # 平均绝对百分比误差
        if 'mape' in metrics:
            results['mape'] = np.mean(np.abs((actual_filtered - predicted_filtered) / actual_filtered)) * 100
        
        # 加权绝对百分比误差
        if 'wape' in metrics:
            results['wape'] = (np.sum(np.abs(actual_filtered - predicted_filtered)) / np.sum(actual_filtered)) * 100
        
        return results
    
    def generate_safety_stock(self, demand_data, service_level=0.95, lead_time_days=7):
        """
        计算安全库存
        
        Args:
            demand_data: 历史需求数据
            service_level: 服务水平（0-1）
            lead_time_days: 采购提前期（天）
            
        Returns:
            dict: 安全库存计算结果
        """
        if len(demand_data) < 30:
            return {'safety_stock': 0, 'reorder_point': 0, 'max_stock': 0}
        
        # 计算日需求统计
        daily_demand = demand_data['quantity']
        avg_daily_demand = daily_demand.mean()
        std_daily_demand = daily_demand.std()
        
        # 计算提前期需求统计
        lead_time_demand = avg_daily_demand * lead_time_days
        lead_time_std = std_daily_demand * np.sqrt(lead_time_days)
        
        # 根据服务水平计算Z值
        from scipy.stats import norm
        z_score = norm.ppf(service_level)
        
        # 计算安全库存
        safety_stock = z_score * lead_time_std
        
        # 计算再订货点
        reorder_point = lead_time_demand + safety_stock
        
        # 计算最大库存
        max_stock = reorder_point + avg_daily_demand * 7  # 额外7天的需求
        
        return {
            'avg_daily_demand': avg_daily_demand,
            'std_daily_demand': std_daily_demand,
            'safety_stock': max(0, safety_stock),
            'reorder_point': max(0, reorder_point),
            'max_stock': max(0, max_stock),
            'service_level': service_level,
            'lead_time_days': lead_time_days
        }
    
    def forecast_all_skus(self, demand_data_dict, method='ensemble', horizon=30, 
                         min_history_days=14):
        """
        批量预测所有SKU
        
        Args:
            demand_data_dict: 各SKU的需求数据字典
            method: 预测方法
            horizon: 预测期长度
            min_history_days: 最小历史数据要求
            
        Returns:
            dict: 各SKU的预测结果
        """
        forecasts = {}
        
        for sku, data in demand_data_dict.items():
            if len(data) < min_history_days:
                print(f"SKU {sku} 历史数据不足({len(data)}天)，跳过预测")
                continue
            
            try:
                if method == 'ensemble':
                    forecast = self.ensemble_forecast(data, horizon=horizon)
                elif method in self.forecast_methods:
                    forecast = self.forecast_methods[method](data, horizon=horizon)
                else:
                    forecast = self._moving_average_forecast(data, horizon=horizon)
                
                forecasts[sku] = {
                    'forecast': forecast,
                    'method': method,
                    'history_length': len(data),
                    'safety_stock': self.generate_safety_stock(data)
                }
                
            except Exception as e:
                print(f"SKU {sku} 预测失败: {str(e)}")
                continue
        
        return forecasts

class ForecastValidator:
    """预测验证器"""
    
    @staticmethod
    def validate_forecast_quality(forecasts, actual_data):
        """验证预测质量"""
        validation_results = {}
        
        for sku, forecast_data in forecasts.items():
            if sku not in actual_data:
                continue
            
            predicted = forecast_data['forecast']
            actual = actual_data[sku]
            
            # 对齐时间索引
            common_index = predicted.index.intersection(actual.index)
            if len(common_index) == 0:
                continue
            
            predicted_aligned = predicted.loc[common_index]
            actual_aligned = actual.loc[common_index, 'quantity']
            
            # 计算准确性指标
            forecaster = DemandForecaster()
            accuracy = forecaster.calculate_forecast_accuracy(
                actual_aligned.values, 
                predicted_aligned.values
            )
            
            validation_results[sku] = {
                'accuracy_metrics': accuracy,
                'common_points': len(common_index),
                'method': forecast_data['method']
            }
        
        return validation_results
    
    @staticmethod
    def identify_forecast_bias(validation_results):
        """识别系统性预测偏差"""
        bias_analysis = {}
        
        for sku, results in validation_results.items():
            mape = results['accuracy_metrics'].get('mape', np.inf)
            
            if mape < 10:
                bias_level = '优秀'
            elif mape < 20:
                bias_level = '良好'
            elif mape < 30:
                bias_level = '一般'
            else:
                bias_level = '需改进'
            
            bias_analysis[sku] = {
                'mape': mape,
                'bias_level': bias_level,
                'recommendation': ForecastValidator._get_recommendation(mape)
            }
        
        return bias_analysis
    
    @staticmethod
    def _get_recommendation(mape):
        """根据MAPE提供改进建议"""
        if mape < 10:
            return "预测准确，可微调模型参数"
        elif mape < 20:
            return "建议增加历史数据或调整季节性参数"
        elif mape < 30:
            return "建议检查数据质量，考虑外部因素影响"
        else:
            return "需要重新评估预测方法，考虑业务变化因素"